from hope import loadGraphFromEdgeListTxt
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from connect import egress,ingress,two
from spacetrans import space
from graph import graph

def toNum(fn):
  dG = loadGraphFromEdgeListTxt(fn, directed=True)
  pos = {}
  i = 0
  for n in sorted(dG.nodes()):
    pos[n] = i
    i = i + 1
  g = {}
  with open(fn, 'r') as f:
    for line in f:
      edge = line.strip().split()
      key = (pos[edge[0]], pos[edge[1]])
      if g.__contains__(key):
        g[key] = g[key] + 1
      else:
        g[key] = 1

  rG = nx.DiGraph()
  for k in g:
    rG.add_edge(k[0], k[1], weight=g[k])
  return rG,pos
def ta(x,a):
  x[x>a]=x[x>a]+a
  x[x<a]=x[x<a]*2.
  return x
  
def hole(r,p,a):
  r = r - p
  r = ta(np.absolute(r),a)*np.exp(1j*np.angle(r))
  return r + p

def nodes():
  rG,p = toNum("../match.g")
  eg = egress(rG)
  ig = ingress(rG)
  eg.dump(5)
  ig.dump(5)

  r = initR(rG)
  aa = 3
  def oneStep(r):
    r = ig.oneStep(r)
    r = ig.oneStep(r)
    r = ig.oneStep(r)
    r = eg.oneStep(r)
    return r    
  for i in range(100):
    if (i%20) == 0:
      plot(r)
    center = np.average(r)
    r = r - center
    r = rol(r,aa)
    if (i%20) == 0:
      plot(r)
    c,ra = find(r,12)
    r = hole(r,c,ra*aa)
    if (i%20) == 0:
      plot(r)
    r = oneStep(r)
    if (i%20) == 0:
      plot(r)
  plot(r)   
  for i in range(10):
    r = oneStep(r)
  print("end")
  plot(r)   
  nplot(r,rG)
def limit(r,a):
  n = np.absolute(r)
  r[n>a] = r[n>a]/np.absolute(r[n>a])*a
  return r
def find(r,N):
  n = np.round((np.angle(r)+np.pi)*N/(2*np.pi)).astype(np.int)
  binc = np.bincount(n)
  mmax = np.argmax(binc)
  mmin = np.argmin(binc)
  print("m=",mmax,"binc=",binc)
  mm = binc[mmax]
  ra = mm / (sum(binc)/N)
  if ra>1.2:
    c = np.sum((n==mmax)*r)/mm
  else:
    c = 0+0j
  print("c=",c,"ra",ra)
  return c,ra/2
def findZ(r,N):
  nr = np.round(r.real+N/2-0.5).astype(np.int)
  nr[nr>N] = N
  ni = np.round(r.imag+N/2-0.5).astype(np.int)
  n = nr*N+ni
  n[n<0] = 0
  n[n>N*N] = N*N
  binc = np.bincount(n)
  mmax = np.argmax(binc)
  mmin = np.argmin(binc)
  # print("m=",mmax,"binc=",binc)
  mm = binc[mmax]
  ra = mm / (np.sum(binc)/(N*N))
  if ra>(4):
    c = np.sum((n==mmax)*r)/mm
  else:
    c = 0+0j
  print("c=",c,"ra",ra)
  return c,ra/N
def init(G):
  node_num = nx.number_of_nodes(G)
  r = np.exp(-1j*2*np.pi*np.arange(node_num)/float(node_num))
  return r

def initR(G):
  node_num = nx.number_of_nodes(G)
  r = fromR(np.arange(node_num)/10.,np.zeros((node_num,),dtype=np.float),5)
  return r
def trol(r):
  l,=r.shape
  rr = np.sum(r.real*r.real)/l
  ii = np.sum(r.imag*r.imag)/l
  ri = np.sum(r.real*r.imag)/l
  return rr,ii,ri
def rol(r,aa):
  rr,ii,ri = trol(r)
  # print("1",rr,ii,ri)
  a = np.array([[rr,ri],[ri,ii]])
  D,L = np.linalg.eigh(a)
  # print("invL",L,D)
  st = L[0,0]-1j*L[0,1]
  vr = r * st
  vr = (vr.real/np.sqrt(D[0])+1j*vr.imag/np.sqrt(D[1]))*aa
  # rr,ii,ri = trol(vr)
  # print("2",rr,ii,ri)
  # plot(vr)
  return vr
# def oneStep(pos,nb,w):
#   # pos = pos * (1+0.03j)
#   r = np.zeros(pos.shape,dtype=np.complex128)
#   for i in range(len(pos)):
#     r[i] = np.average(pos[nb[i]],weights=w[i])
#   return r

def plot(r):
  plt.scatter(r.real,r.imag,s=0.1)
  plt.show()

def nplot(r,G):
  pos = {}
  for n in G.nodes():
    pos[n] = np.array([r[n].real,r[n].imag])
  nx.draw_networkx(G, pos, labels=None)
  plt.show()

# def neighbors(G):
#   r = []
#   w = []
#   for n in range(nx.number_of_nodes(G)):
#     r0 = [n]
#     w0 = [1]
#     for nb in nx.all_neighbors(G,n):
#       r0.append(nb)
#       weight = G.get_edge_data(nb,n)
#       if weight==None:
#         ww = 0
#       else:
#         ww = weight['weight']
#       w0.append(ww)
#     r.append(r0)
#     w.append(w0)
#   return r,w

def toR(r):
  a = np.absolute(r)
  p = np.angle(r)
  l = np.floor(a)*2*np.pi+p
  da = a - l/(2*np.pi)
  return l,da
def fromR(p,da,N):
  index = np.argsort(p)
  sizel, = p.shape
  l = int(np.ceil(sizel / (N*N)))
  r = np.zeros(p.shape,dtype=np.complex128)
  for i in range(N):
    pl = (2*i+1)*l
    ps = i*i*l
    for j in range(pl):
      k = ps + j
      if k<sizel:
        npp = (i+float(j)/float(pl))
        na = npp*2*np.pi
        npda = npp+da[index[k]]
        r[index[k]] = npda * np.exp(1j*na)
  return r

def main():
  gp = graph("../test_flatten.g")
  # gp = graph("../match.g")
  rG = gp.nG
  eg = egress(rG)
  ig = ingress(rG)
  tw = two(rG)
  eg.dump(5)
  ig.dump(5)
  tw.dump(5)

  def oneStep(r):
    r = ig.oneStep(r)
    r = eg.oneStep(r)
    return r 
  
  node_num = nx.number_of_nodes(rG)
  sp = gp.fromEMBFile("../test_flatten.emb")
  # sp.init(node_num)

  # sp.oneStep(ig)
  # print("3",gp.power(sp.pos))
  # sp.up(eg,ig)
  # print("setp 2")
  
  # sp.oneStep(ig)
  sp.oneStep(ig)
  sp.oneStep(eg)
  sp.down(2)

  # sp.oneStep(ig)
  # sp.oneStep(ig)
  # sp.oneStep(eg)

  # sp.oneStep(ig)
  # sp.oneStep(ig)
  # sp.oneStep(eg)

  
  # print("setp 3")
  # print("4",gp.power(sp.pos))
  
  aa = 3
  
  r = sp.toC()*aa
  
  for i in range(5):
    center = np.average(r)
    r = r - center
    r = rol(r,aa)
    print(i)
    # plot(r)
    c,ra = find(r,24)
    r = hole(r,c,1.4)
    for k in range(3):
      c,ra = findZ(r,24)
      r = hole(r,c,1.4)
      
      print("find",k)
    for k in range(3):
      r = oneStep(r)
      sp.fromC(r)
      print("*",gp.power(sp.pos))
  
  sp.fromC(r)
  
  print("5",gp.power(sp.pos))
  sp.up(eg,ig)
  sp.oneStep(ig)
  sp.oneStep(eg)
  sp.up(eg,ig)
  sp.oneStep(ig)
  sp.oneStep(eg)
  sp.down(2)
  print("6",gp.power(sp.pos))
  sp.plot()
  for i in range(10):
    sp.up(eg,ig)
    sp.oneStep(ig)
    sp.oneStep(eg)
    sp.up(eg,ig)
    sp.oneStep(ig)
    sp.oneStep(eg)
    sp.down(2)
    r = sp.toC()
    sp.fromC(r)
    print(i,gp.power(sp.pos))
    sp.plot()
  sp.nplot(rG)

def test():
  gp = graph("../test_flatten.g")
  # gp = graph("../match.g")
  rG = gp.nG
  eg = egress(rG)
  ig = ingress(rG)
  eg.dump(5)
  ig.dump(5)

  node_num = nx.number_of_nodes(rG)
  sp = space()
  sp.init(node_num)
  
  print("0",gp.power(sp.pos))
  sp.up(eg,ig)
  print("setp 1")
  sp.oneStep(ig)
  print("1",gp.power(sp.pos))
  sp.oneStep(ig)
  print("2",gp.power(sp.pos))
  sp.oneStep(eg)
  print("3",gp.power(sp.pos))
  sp.up(eg,ig)
  print("setp 2")
  sp.oneStep(ig)
  sp.oneStep(ig)
  sp.oneStep(eg)

  sp.down(2)
  print("setp 3")
  print("4",gp.power(sp.pos))  

  sp.plot()
if __name__ == '__main__':
  main()
  